import numpy as np

from .molecules import *
from .gto_integral import *
from .electron_cloud import *

# Bohr radius
a0 = 0.52917721067

# 1 hartree energy = 27.211386245988eV
hartree = 27.211386245988


def HartreeFock(mol: Molecule, delta_E=1e-10, max_iter=100, details=False, verbose=0) -> float:
    '''
    Hartree-Fock自洽场计算
    参数：
        mol:      待计算的分子
        delta_E:  收敛能量（两次计算的能量差小于此能量即停止计算）
        max_iter: 最大迭代次数，一般来说不会达到100次迭代
        details:  是否返回轨道能量，基矢和波函数展开系数
        verbose:  打印信息的等级
            0: 不打印信息
            1: 打印每次迭代的电子能量
    返回值：
        (E, ei, basis, Cx)
        E:     基态能量
        ei:    单粒子轨道能量
        basis: 基矢
        Cx:    波函数展开系数
        如果 details == False，则只返回基态能量
    '''
    basis = []        # 获取分子轨道
    Ne = -mol.charge  # 获取分子总电子数
    for atom in mol.atoms:
        Ne += atom.Z
        basis.append(atom.orbital)

    assert Ne % 2 == 0, "目前暂不支持奇数电子数的计算"

    Nb = len(basis)     # 轨道数
    Nn = len(mol.atoms)  # 原子数

    # 重叠积分
    S = np.zeros((Nb, Nb))
    for i in range(Nb):
        for j in range(i, Nb):
            S[i, j] = S[j, i] = overlap_integral_sto3g(basis[i], basis[j])

    # 动能积分
    T = np.zeros((Nb, Nb))
    for i in range(Nb):
        for j in range(i, Nb):
            T[i, j] = T[j, i] = kinetic_integral_sto3g(basis[i], basis[j])

    # 原子核势能
    nuclear_energy = 0.0
    for i in range(Nn):
        for j in range(i+1, Nn):
            nuclear_energy += mol.atoms[i].Z * mol.atoms[j].Z / \
                abs(mol.atoms[i].R - mol.atoms[j].R)

    # 原子核吸引势积分
    Vn = np.zeros((Nb, Nb))
    for atom in mol.atoms:
        for i in range(Nb):
            for j in range(i, Nb):
                _v = nuclear_attraction_integral_sto3g(
                    atom.Z, atom.R, basis[i], basis[j])
                Vn[i, j] += _v
                if i != j:
                    Vn[j, i] += _v

    # 单体矩阵元
    H0 = T + Vn

    # 电子排斥积分
    Ve = np.zeros((Nb, Nb, Nb, Nb))
    for i in range(Nb):
        for j in range(Nb):
            for k in range(Nb):
                for l in range(Nb):
                    Ve[i, j, k, l] = electron_repulsion_integral_sto3g(
                        basis[i], basis[j], basis[k], basis[l]
                    )

    d, U = np.linalg.eigh(S)
    X = U @ np.diag(d ** -0.5) @ np.transpose(U)

    electron_energy = 0.  # 电子能量
    old_energy = 0.       # 上一次计算的电子能量
    iter_time = 0

    rho = np.zeros((Nb, Nb))   # 密度矩阵
    G = np.zeros((Nb, Nb))   # 有效相互作用矩阵
    F = np.zeros((Nb, Nb))   # Fock 矩阵
    H = np.zeros((Nb, Nb))   # 对角基下的哈密顿量
    C = np.zeros((Nb, Nb))   # 对角基下的展开系数
    Cx = np.zeros((Nb, Nb))   # 原基下的展开系数
    ei = np.zeros((Nb,))      # 轨道能量

    for hf_iter in range(max_iter):
        F = H0 + G
        H = X @ F @ X
        ei, C = np.linalg.eigh(H)
        Cx = np.real(X @ C)

        # 更新密度矩阵
        rho[:] = 0
        for i in range(Ne//2):
            for a in range(Nb):
                for b in range(Nb):
                    rho[a, b] += 2 * Cx[a, i] * Cx[b, i]

        # 更新能量
        electron_energy = 0.0
        for a in range(Nb):
            for b in range(Nb):
                electron_energy += rho[b, a] * (H0[a, b] + F[a, b])

        electron_energy *= 0.5

        if verbose == 1:
            print(f'electron energy = {electron_energy:.12f} a.u.')

        if abs(old_energy - electron_energy) < delta_E:
            break

        old_energy = electron_energy

        # 更新有效相互作用
        G[:] = 0
        for a in range(Nb):
            for b in range(Nb):
                for c in range(Nb):
                    for d in range(Nb):
                        G[a, d] += (Ve[a, b, c, d] - 0.5 *
                                    Ve[a, b, d, c]) * rho[c, b]

        iter_time = hf_iter

    total_energy = electron_energy + nuclear_energy
    if iter_time == max_iter:
        print("Warning: Hartree-Fock计算达到了最大迭代次数，可能还没有收敛")

    if not details:
        return total_energy
    else:
        return total_energy, ei, basis, Cx
